# Copyright 2025 AlQuraishi Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Linear layer initialization configuration for AlphaFold3 model."""

from ml_collections import ConfigDict

########################
# Primitives
########################

# AF3
swiglu_init = ConfigDict(
    {
        "linear_a": {"bias": False, "init": "relu"},
        "linear_b": {"bias": False, "init": "relu"},
    }
)

# AF3
ada_ln_init = ConfigDict(
    {
        "linear_g": {"bias": True, "init": "final"},
        "linear_s": {"bias": False, "init": "final"},
    }
)

mha_init = ConfigDict(
    {
        "linear_q": {"bias": False, "init": "default"},
        "linear_k": {"bias": False, "init": "default"},
        "linear_v": {"bias": False, "init": "default"},
        "linear_g": {"bias": False, "init": "gating"},
        "linear_o": {"bias": False, "init": "final"},
    }
)

att_pair_bias_mha_init = ConfigDict(
    {
        "linear_q": {"bias": True, "init": "default"},
        "linear_k": {"bias": False, "init": "default"},
        "linear_v": {"bias": False, "init": "default"},
        "linear_g": {"bias": False, "init": "gating"},
        "linear_o": {"bias": False, "init": "final"},
    }
)

att_pair_bias_mha_ada_init = ConfigDict(
    {
        "linear_q": {"bias": True, "init": "default"},
        "linear_k": {"bias": False, "init": "default"},
        "linear_v": {"bias": False, "init": "default"},
        "linear_g": {"bias": False, "init": "gating"},
        "linear_o": {"bias": False, "init": "default"},
    }
)

mha_bias_init = ConfigDict(
    {
        "linear_z": {"bias": False, "init": "default"},
        "mha": mha_init,
    }
)

########################
# Layers
########################

att_pair_bias_init = ConfigDict(
    {
        "linear_z": {"bias": False, "init": "default"},
        "layer_norm_z": {"create_offset": True},
        "mha": att_pair_bias_mha_init,
    }
)

diffusion_att_pair_bias_init = ConfigDict(
    {
        "ada_ln": ada_ln_init,
        "linear_ada_out": {"bias": True, "init": "gating_ada_zero"},
        "linear_z": {"bias": False, "init": "default"},
        "layer_norm_z": {"create_offset": False},
        "mha": att_pair_bias_mha_ada_init,
    }
)

tri_mul_init = ConfigDict(
    {
        "linear_a_p": {"bias": False, "init": "default"},
        "linear_a_g": {"bias": False, "init": "gating"},
        "linear_b_p": {"bias": False, "init": "default"},
        "linear_b_g": {"bias": False, "init": "gating"},
        "linear_g": {"bias": False, "init": "gating"},
        "linear_z": {"bias": False, "init": "final"},
    }
)

# Not used by default, but config is included in case the
# "fuse_projection_weights" option is set to True
fused_tri_mul_init = ConfigDict(
    {
        "linear_ab_p": {"bias": False, "init": "default"},
        "linear_ab_g": {"bias": False, "init": "gating"},
        "linear_g": {"bias": False, "init": "gating"},
        "linear_z": {"bias": False, "init": "final"},
    }
)

tri_att_init = mha_bias_init

opm_init = ConfigDict(
    {
        "linear_1": {"bias": False, "init": "default"},
        "linear_2": {"bias": False, "init": "default"},
        "linear_out": {"bias": True, "init": "final"},
    }
)

msa_pair_avg_init = ConfigDict(
    {
        "linear_z": {"bias": False, "init": "default"},
        "linear_v": {"bias": False, "init": "default"},
        "linear_g": {"bias": False, "init": "gating"},
        "linear_o": {"bias": False, "init": "final"},
    }
)

swiglu_transition_init = ConfigDict(
    {
        "swiglu": swiglu_init,
        "linear_out": {"bias": False, "init": "final"},
    }
)

relu_transition_init = ConfigDict(
    {
        "layers": {"bias": True, "init": "relu"},
        "linear_out": {"bias": True, "init": "final"},
    }
)

transition_init = {"swiglu": swiglu_transition_init, "relu": relu_transition_init}

cond_transition_init = ConfigDict(
    {
        "ada_ln": ada_ln_init,
        "swiglu": swiglu_init,
        "linear_g": {"bias": True, "init": "gating_ada_zero"},
        "linear_out": {"bias": False, "init": "default"},
    }
)

diffusion_transformer_init = ConfigDict(
    {
        "att_pair_bias": diffusion_att_pair_bias_init,
        "cond_transition": cond_transition_init,
    }
)

# AF3
ref_atom_emb_init = ConfigDict(
    {
        "linear_feats": {"bias": False, "init": "default"},
        "linear_ref_offset": {"bias": False, "init": "default"},
        "linear_inv_sq_dists": {"bias": False, "init": "default"},
        "linear_valid_mask": {"bias": False, "init": "default"},
    }
)

noisy_pos_emb_init = ConfigDict(
    {
        "linear_s": {"bias": False, "init": "final"},
        "linear_z": {"bias": False, "init": "final"},
        "linear_r": {"bias": False, "init": "default"},
    }
)

atom_att_enc_init = ConfigDict(
    {
        "ref_atom_emb": ref_atom_emb_init,
        "noisy_pos_emb": noisy_pos_emb_init,
        "linear_l": {"bias": False, "init": "default"},
        "linear_m": {"bias": False, "init": "default"},
        "pair_mlp_1": {"bias": False, "init": "relu"},
        "pair_mlp_2": {"bias": False, "init": "relu"},
        "pair_mlp_3": {"bias": False, "init": "final"},
        "diffusion_transformer": diffusion_transformer_init,
        "linear_q": {"bias": False, "init": "default"},
    }
)

atom_att_dec_init = ConfigDict(
    {
        "linear_q_in": {"bias": False, "init": "default"},
        "diffusion_transformer": diffusion_transformer_init,
        "linear_q_out": {"bias": False, "init": "final"},
    }
)

diffusion_cond_init = ConfigDict(
    {
        "linear_z": {"bias": False, "init": "default"},
        "transition_z": swiglu_transition_init,
        "linear_s": {"bias": False, "init": "default"},
        "linear_n": {"bias": False, "init": "default"},
        "transition_s": swiglu_transition_init,
    }
)

########################
# Feature Embedders
########################

input_emb_init = ConfigDict(
    {
        "linear_s": {"bias": False, "init": "default"},
        "linear_z_i": {"bias": False, "init": "default"},
        "linear_z_j": {"bias": False, "init": "default"},
        "linear_relpos": {"bias": False, "init": "default"},
        "linear_token_bonds": {"bias": False, "init": "default"},
    }
)

msa_module_emb_init = ConfigDict(
    {
        "linear_m": {"bias": False, "init": "default"},
        "linear_s_input": {"bias": False, "init": "default"},
    }
)

# TODO: check initialization
templ_pair_feat_emb_init = ConfigDict(
    {
        "linear_a": {"bias": False, "init": "default"},
        "linear_z": {"bias": False, "init": "relu"},
    }
)

########################
# Heads
########################

pairformer_head_init = ConfigDict(
    {
        "linear_i": {"bias": False, "init": "default"},
        "linear_j": {"bias": False, "init": "default"},
        "linear_distance": {"bias": False, "init": "default"},
    }
)

pae_init = ConfigDict({"linear": {"bias": False, "init": "final"}})

pde_init = ConfigDict({"linear": {"bias": False, "init": "final"}})

lddt_init = ConfigDict({"linear": {"bias": False, "init": "final"}})

exp_res_all_atom_init = ConfigDict({"linear": {"bias": False, "init": "final"}})

distogram_init = ConfigDict({"linear": {"bias": False, "init": "final"}})

########################
# Latent
########################

pair_block_init = ConfigDict(
    {
        "tri_mul": tri_mul_init,
        "fused_tri_mul": fused_tri_mul_init,
        "tri_att": tri_att_init,
        "pair_transition": transition_init,
    }
)

msa_block_init = ConfigDict(
    {
        "msa_row_att": mha_bias_init,
        "msa_transition": transition_init,
        "opm": opm_init,
        "pair_block": pair_block_init,
    }
)

msa_module_init = ConfigDict(
    {
        **msa_block_init,
        "msa_pair_avg": msa_pair_avg_init,
    }
)

pairformer_init = ConfigDict(
    {
        "pair_block": pair_block_init,
        "att_pair_bias": att_pair_bias_init,
        "transition": swiglu_transition_init,
    }
)

templ_module_init = ConfigDict({"linear_t": {"bias": False, "init": "default"}})

########################
# Structure
########################

# TODO: Maybe structure this like other modules, where the configs are contained
# in one dict for the full module. Because of the way the params are passed,
# we only need to define the top level linear layers here.
diffusion_module_init = ConfigDict(
    {
        "linear_s": {"bias": False, "init": "final"},
    }
)
